import numpy as np
from scipy.integrate import solve_ivp
import matplotlib.pyplot as plt
import matplotlib

# 设置中文字体（解决中文显示问题）
plt.rcParams['font.sans-serif'] = ['SimHei']  # 使用黑体
plt.rcParams['axes.unicode_minus'] = False  # 解决负号显示问题


# 单摆运动方程（无量纲形式）
def pendulum_eq(t, y, beta, f, Omega):
    theta, omega = y
    dtheta_dt = omega
    domega_dt = -2 * beta * omega - np.sin(theta) + f * np.cos(Omega * t)
    return [dtheta_dt, domega_dt]


# 模拟参数设置
beta = 0.0  # 无量纲阻尼系数 0.25
f = 1.35  # 无量纲驱动力振幅 (尝试1.07, 1.15, 1.35, 1.45, 1.47, 1.50)
Omega = 2 / 3  # 无量纲驱动力频率

# 初始条件 [theta, omega]
y0 = [0.1, 0.0]  # 小角度初始位移

# 时间范围 (无量纲时间，相当于ω₀t)
t_start = 0
t_end = 300  # 增加总时间以更好地观察稳态行为
t_points = 6000  # 增加时间点数以提高精度
t_span = (t_start, t_end)
t_eval = np.linspace(t_start, t_end, t_points)

# 求解微分方程
sol = solve_ivp(pendulum_eq, t_span, y0, args=(beta, f, Omega),
                t_eval=t_eval, rtol=1e-8, atol=1e-8)

# 提取结果
theta = sol.y[0]
omega = sol.y[1]
t = sol.t

# 绘制θ-t图
plt.figure(figsize=(12, 9))
plt.subplot(2, 1, 1)
plt.plot(t, theta, 'b-', linewidth=1)
plt.title(f'受驱阻尼单摆运动 (β={beta}, f={f}, Ω={Omega:.3f})', fontsize=14)
plt.xlabel('无量纲时间 (ω₀t)', fontsize=12)
plt.ylabel('摆角 θ (弧度)', fontsize=12)
plt.grid(alpha=0.3)

# 绘制相图 (θ-ω图)
plt.subplot(2, 1, 2)
plt.plot(theta, omega, 'g-', linewidth=0.5)
plt.title('相位空间图', fontsize=14)
plt.xlabel('摆角 θ (弧度)', fontsize=12)
plt.ylabel('角速度 ω (dθ/dt)', fontsize=12)
plt.grid(alpha=0.3)

plt.tight_layout()
plt.show()

# 绘制庞加莱截面图（高效方法）
# if f >= 1.15:  # 仅在混沌参数下绘制
#     drive_period = 2 * np.pi / Omega
#
#     # 从现有解中提取庞加莱截面点
#     # 跳过前50个驱动周期对应的数据点
#     start_index = int(50 * drive_period / (t_end / t_points))
#     poincare_indices = []
#
#     # 查找每个驱动周期结束时的数据点
#     for i in range(start_index, len(t)):
#         # 检查是否接近驱动周期的整数倍
#         if abs((t[i] % drive_period) - 0) < 0.1:
#             poincare_indices.append(i)
#
#     if len(poincare_indices) > 0:
#         poincare_theta = theta[poincare_indices]
#         poincare_omega = omega[poincare_indices]
#
#         plt.figure(figsize=(8, 6))
#         plt.plot(poincare_theta, poincare_omega, 'ro', markersize=2)
#         plt.title(f'庞加莱截面 (β={beta}, f={f}, Ω={Omega:.3f})', fontsize=14)
#         plt.xlabel('摆角 θ (弧度)', fontsize=12)
#         plt.ylabel('角速度 ω (dθ/dt)', fontsize=12)
#         plt.grid(alpha=0.3)
#         plt.tight_layout()
#         plt.show()
#     else:
#         print(f"警告: 无法提取庞加莱截面点")
